package com.sunzm.spark.streaming.kafka

import org.apache.commons.lang3.time.DateFormatUtils
import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord}
import org.apache.kafka.common.serialization.StringDeserializer
import org.apache.spark.SparkConf
import org.apache.spark.streaming.dstream.{DStream, InputDStream}
import org.apache.spark.streaming.kafka010.ConsumerStrategies.Subscribe
import org.apache.spark.streaming.kafka010.{KafkaUtils, LocationStrategies}
import org.apache.spark.streaming.{Seconds, StreamingContext}

object DStreamKafkaDemo {
  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf()
      .setAppName(s"${this.getClass.getSimpleName.stripSuffix("$")}")
      .setMaster("local[*]")
    val ssc = new StreamingContext(sparkConf, Seconds(5))

    val kafkaParams = Map[String, Object](
      ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG -> "82.156.210.70:9093",
      ConsumerConfig.GROUP_ID_CONFIG -> "use_a_separate_group_id_for_each_stream",
      ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG -> classOf[StringDeserializer],
      ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG -> classOf[StringDeserializer])

    // bin/kafka-console-producer.sh --bootstrap-server 10.0.8.11:9092 --topic my-topic
    val topics = Array("my-topic")
    val kafkaDStream: InputDStream[ConsumerRecord[String, String]] = KafkaUtils.createDirectStream[String, String](
      ssc,
      LocationStrategies.PreferConsistent,
      Subscribe[String, String](topics, kafkaParams)
    )

    val wordCounts: DStream[(String, Int)] = kafkaDStream.map(_.value).transform(rdd => {

      rdd.flatMap(_.split(","))
        .map((_, 1))
        .reduceByKey(_ + _)

    })

    wordCounts.foreachRDD((rdd, time) => {

      val milliseconds = time.milliseconds

      val timeStr = DateFormatUtils.format(milliseconds, "yyyy-MM-dd HH:mm:ss")

      if (!rdd.isEmpty()) {

        val partitionsLength = rdd.partitions.length

        println(s"当前计算时间: ${timeStr}, 当前RDD的分区数： ${partitionsLength}")

        rdd.foreach(line => {
          println(s"单词个数统计结果: ${line}")
        })
      } else {
        println(s"没有需要计算的数据: ${timeStr}")
      }

    })

    ssc.start()
    ssc.awaitTermination()
  }
}
